
from pyglet_render.viewer import Viewer
from ship_envs.core.ship_forward import ShipForward
from ship_envs.core.model_transformer import model2geom
from ship_envs.environment.single_ship_env import SingleShipEnv


class SingleShipEnvS1(SingleShipEnv):
    def __init__(self):
        super().__init__()
        self.world.add_ship(ShipForward())

    def reset(self):
        # 重置船位置及目标点位置
        ship = self.get_ship()
        ship.cur_speed = 0.004
        pos = (0.2, 0.5)
        tar_pos = (0.8, 0.5)
        ship.set_pos(pos)
        ship.target.set_pos(tar_pos)

        other_ship = self.world.other_ships[0]
        pos = (0.5, 0.2)
        angle = 90
        other_ship.set_pos(pos)
        other_ship.set_angle(angle)

        # 重置步数
        self.step_count = 0
        # 刷新obs
        self.world.refresh()
        # 单智能体，obs为整个world的obs  全观测
        return self.observation_callback(self.world, ship)


